library(lpSolve)

load("nelson_data.RData")
attach(kkk)

# Dichotomize variables so they match definitions in the paper
Y <- ifelse(kspeech > median(kspeech), 1, 0)
X <- vidcdum
M <- ifelse(idisord > median(idisord), 1, 0)

# Observed strata
P000 <- sum(Y==0 & M==0 & X==0)/sum(X==0)
P001 <- sum(Y==0 & M==0 & X==1)/sum(X==1)
P010 <- sum(Y==0 & M==1 & X==0)/sum(X==0)
P011 <- sum(Y==0 & M==1 & X==1)/sum(X==1)
P100 <- sum(Y==1 & M==0 & X==0)/sum(X==0)
P101 <- sum(Y==1 & M==0 & X==1)/sum(X==1)
P110 <- sum(Y==1 & M==1 & X==0)/sum(X==0)
P111 <- sum(Y==1 & M==1 & X==1)/sum(X==1)
P <- c(P000, P001, P010, P011, P100, P101, P110, P111)

# Create the constraint (A) matrix for lp
A.obs <- matrix(c(
        rep(c(1,1,0,0), 2^3), rep(0, 2^5), #P000
        rep(c(rep(c(1,0),2^2), rep(0,2^3)), 2^2), #P001
        rep(c(rep(c(0,0,1,1), 2^2), rep(0, 2^4)), 2), #P010
        rep(c(0,1,0,1,0,0,0,0), 2^3), #P011
        rep(0, 2^5), rep(c(1,1,0,0), 2^3), #P100
        rep(c(rep(0,2^3), rep(c(1,0), 2^2)), 2^2), #P101
        rep(c(rep(0, 2^4), rep(c(0,0,1,1), 2^2)), 2), #P010
        rep(c(0,0,0,0,0,1,0,1), 2^3), #P011
        rep(1, 2^6) # sum=1
        ), nrow=9, byrow=TRUE)

A00 <- 1/(P000 + P100)
A01 <- 1/(P001 + P101)
A10 <- 1/(P010 + P110)
A11 <- 1/(P011 + P111)

A.1 <- diag(1,2^4) %x% matrix(c(-A01,A11,-A01,A11), nrow=1)
A.0 <- diag(1,2^4) %x% matrix(c(-A00,-A00,A10,A10), nrow=1)

f.con <- rbind(A.obs,A.1,A.0,A.1,A.0)

# Set up rho & result container
rho <- round(seq(0, 1, by = .001),3)
bounds <- matrix(NA, ncol=length(rho), nrow=5)
colnames(bounds) <- rho
rownames(bounds) <- c("delta.1.up", "delta.1.lo", "delta.0.up", "delta.0.lo", "feasible")

# Loop w.r.t rho
for(i in seq(along=rho)){
    f.dir <- c(rep("=", 9), rep("<=", 32), rep(">=", 32))
    f.rhs <- c(P,1,rep(rho[i], 32), rep(-rho[i],32))
    f.obj.1 <- rep(c(0,0,0,0,0,1,-1,0,0,-1,1,0,0,0,0,0), 2^2)
    f.obj.0 <- c(rep(0, 2^4), rep(c(0,1,-1,0), 2^2), rep(c(0,-1,1,0), 2^2), rep(0,2^4))
    res.1.max <- lp("max", f.obj.1, f.con, f.dir, f.rhs)
    res.1.min <- lp("min", f.obj.1, f.con, f.dir, f.rhs)
    res.0.max <- lp("max", f.obj.0, f.con, f.dir, f.rhs)
    res.0.min <- lp("min", f.obj.0, f.con, f.dir, f.rhs)
    bounds[1,i] <- ifelse(res.1.max$status == 0, res.1.max$objval, NA)
    bounds[2,i] <- ifelse(res.1.min$status == 0, res.1.min$objval, NA)
    bounds[3,i] <- ifelse(res.0.max$status == 0, res.0.max$objval, NA)
    bounds[4,i] <- ifelse(res.0.min$status == 0, res.0.min$objval, NA)
    bounds[5,i] <- ifelse(sum(is.na(bounds[1:4,i]))>0, 0, 1)
}

# Now plot
pdf("nonpara-sens.pdf", height=6.5, width=6.5)
par(mfrow=c(2,1), mar=c(4.1,4.1,2,2))
plot(0,0,type="n", xlim=c(0,1), ylim=c(-1,1), 
    xlab=expression(paste("Sensitivity Parameter: ", upsilon)),
    ylab=expression(paste("Average Mediation Effect: ", bar(delta),"(",1,")")),
    main="")
abline(h=0)
abline(h=bounds[1,1], lty="dashed")
lines(rho, bounds[1,], lwd=2)
lines(rho, bounds[2,], lwd=2)

plot(0,0,type="n", xlim=c(0,1), ylim=c(-1,1), 
    xlab=expression(paste("Sensitivity Parameter: ", upsilon)),
    ylab=expression(paste("Average Mediation Effect: ", bar(delta),"(",0,")")),
    main="")
abline(h=0)
lines(rho, bounds[3,], lwd=2)
lines(rho, bounds[4,], lwd=2)
abline(h=bounds[3,1], lty="dashed")

dev.off()


